﻿module CSSolver
open CSModel
open CSValidation   

module CSP =
    let vdef = function VarDef(n,d) -> n,d
    let vname = vdef>>fst
    let vdomain = vdef>>snd

    let varOrdering = function (VarDef(_,d)) -> 
                                match d with
                                | Ds s -> s |> Seq.length
                                | Di s -> s |> Seq.length

    let tuples (vars:seq<Def>) =
        let nds = vars|>Seq.map vdef |> Seq.toList
        seq {
            let rec loop acc (ls:list<string*Domain>) =
                    match ls with
                    | [] -> seq{yield acc}
                    | (n,Ds d)::rest | (n,Di d)::rest -> 
                        seq {
                            for dv in d do
                            yield! loop (acc|>Map.add n dv) rest
                            }
            yield! loop Map.empty nds
            }

    type ValType = Vi of int | Vs of int | Vb of bool

    let evalTuple exp (tplMap:Map<string,int>) =
        let rec evalNum e = eval e |> function Vi n->n | _ -> failwith "Non numeric value"
        and evalBool e = eval e |> function Vb n->n | _ -> failwith "Non boolean value"
        and eval e =
            match e with
            | Var (VarDef(n,Ds _)) -> Vs tplMap.[n]
            | Var (VarDef(n,Di _)) -> Vi tplMap.[n]
            | Val (ValI i) -> Vi i
            | Val (ValS s) -> Vs s
            | Plus (a,b) -> Vi (evalNum a +  evalNum b)
            | Minus (a,b)-> Vi (evalNum a -  evalNum b)
            | Div (a,b)  -> Vi (evalNum a /  evalNum b)
            | Mod (a,b)  -> Vi (evalNum a %  evalNum b)
            | Mult(a,b)  -> Vi (evalNum a *  evalNum b)
            | Ge (a,b)   -> Vb (evalNum a >= evalNum b)
            | Gt (a,b)   -> Vb (evalNum a >  evalNum b)
            | Le(a,b)    -> Vb (evalNum a <= evalNum b)
            | Lt(a,b)    -> Vb (evalNum a <  evalNum b)
            | Not x      -> Vb (not (evalBool x))
            | Abs x      -> Vi (abs (evalNum x))
            | Eq (a,b)   -> Vb (eval a = eval b)
            | Neq (a,b)  -> Vb (eval a <> eval b)
            | And xs     -> Vb (xs|>Seq.exists(fun x -> eval x = Vb(false)))
            | Or xs      -> Vb (xs|>Seq.exists(fun x -> eval x = Vb(true)))
        eval exp |> function Vb b -> b | _ -> failwith "expression did not evaluate to boolean"

    let join (r1s:seq<Map<string,int>>) (r2s:seq<Map<string,int>>) commonVars =
        (Set.empty, r1s) ||> Seq.fold (fun acc r1 ->
            (acc,r2s)||>Seq.fold(fun acc r2 -> 
                if commonVars|>Seq.forall(fun v-> r1.[v]=r2.[v]) then
                    let r3 = [r1;r2]|>Seq.collect(fun a->a|>Map.toSeq) |> Set.ofSeq
                    acc|>Set.add r3
                else
                    acc))
        |> Seq.map(fun xs -> xs |> Map.ofSeq)

    let projectOut dropVar (rs:seq<Map<string,int>>)  =
        let tuples,constraintMap =
            ((Set.empty,Map.empty),rs) ||> Seq.fold (fun (acc,map) r ->
                let varVal = r.[dropVar]
                let rs = r |> Map.remove dropVar |> Map.toSeq |> Set.ofSeq
                let acc = acc |> Set.add rs
                let map =
                    match map |> Map.tryFind rs with
                    | Some (s) -> map |> Map.add rs (varVal::s)
                    | None -> map |> Map.add rs [varVal]
                acc,map)
        tuples|>Seq.map Map.ofSeq, constraintMap

    type Bucket = 
        {
         VDef:Def
         Exprs:Expr list
         Tuples:(Set<string>*seq<Map<string,int>>) option
         Constraints:Map<Set<string*int>,int list>
         CScope:Set<string>
        }
        static member Default = {Exprs=List.empty; VDef=VarDef("dummy",intSet [1]); Tuples=None; Constraints=Map.empty; CScope=Set.empty}

    let (++) (t1:Map<string,int>) (t2:Map<string,int>) = (t1,t2) ||> Map.fold (fun acc k v -> acc |> Map.add k v)
       
    let variableConsistencySolve (ConstraintSystem(vars,exprs))=
        let orderedVars = vars |> Seq.sortBy varOrdering |> Seq.toList
        let orderedVarNames = orderedVars |> List.map vname

        //initialize buckets
        let bktMap = orderedVars |> Seq.map (fun v->v|>vname,{Bucket.Default with VDef=v})|> Map.ofSeq

        //partition constraints into buckets
        let bktMap = (bktMap,exprs) ||> Seq.fold (fun bmap exp ->
            let vs = exp |> scope |> Seq.map vname |> Set.ofSeq
            let varName = orderedVarNames |> List.find (fun v-> vs|>Set.contains v)
            let bkt = bmap.[varName]
            bmap |> Map.add varName {bkt with Exprs=exp::bkt.Exprs}
            )

        // join and project
        let bktMap = (bktMap,orderedVarNames) ||> Seq.fold(fun bmap varName ->
            let bkt = bmap.[varName]

            let joinedVars,joined = 
                let exprsScope = bkt.Exprs |> Seq.collect scope |> Seq.distinctBy vname 
                match bkt.Tuples with
                | Some (scope,propogatedRelation) ->  
                    let extraVars = exprsScope |> Seq.filter (vname >> scope.Contains >> not)
                    let tuples = tuples extraVars 
                    let js =
                        let joinedTuples = 
                            seq {
                                for p in propogatedRelation do
                                    for t in tuples do
                                        yield p ++ t}
                        (joinedTuples,bkt.Exprs)
                        ||> Seq.fold (fun acc exp -> acc |> Seq.filter (evalTuple exp))
                    let jvs = (scope,exprsScope) ||> Seq.fold (fun acc v -> acc |> Set.add (vname v))
                    jvs,js
                | None ->  
                    if bkt.Exprs |> Seq.isEmpty then
                        Set[bkt.VDef|>vname],tuples [bkt.VDef]
                    else
                        let jvs = exprsScope |> Seq.map vname |> Set.ofSeq
                        let js =
                            (tuples exprsScope,bkt.Exprs)
                            ||> Seq.fold (fun acc exp -> acc |> Seq.filter (evalTuple exp))
                        jvs,js

            let propogationScope = joinedVars |> Set.remove varName
            if propogationScope.IsEmpty then
                let vlist = joined |> Seq.map (fun m->m.[varName]) |> Seq.toList
                let cnstrnts = Map([Set.empty,vlist])
                bmap |> Map.add varName  {bkt with Constraints = cnstrnts; CScope=Set.empty}
            else
                let propogationRelation,constraints = joined |> projectOut varName
                let propogationBucket = orderedVarNames |> List.find (fun n -> propogationScope.Contains n)
            
                let prjBkt = bmap.[propogationBucket]
                let prjBkt = 
                    {prjBkt with
                        Tuples =
                            match prjBkt.Tuples with
                            | Some (prevScope,prevRelation) -> 
                                let commonVars = prevScope |> Set.intersect propogationScope
                                let joined = join prevRelation propogationRelation commonVars
                                let scope = prevScope |> Set.union propogationScope
                                Some(scope,joined)   
                            | None  -> Some(propogationScope, propogationRelation)
                    }
                bmap 
                |> Map.add varName 
                    {bkt with 
                        Constraints=constraints
                        CScope = propogationScope
                    }
                |> Map.add propogationBucket prjBkt
                )
        bktMap,orderedVarNames|>List.rev

type SolutionSet = {Buckets:Map<string,CSP.Bucket>; Variables:string list}
        
let solve css =
    validate css
    let bktmap,varlist = CSP.variableConsistencySolve css
    {Buckets=bktmap;Variables=varlist}

let enumerateSolutions solution =
    seq {
        let rec loop (valsSet:Set<string*int>) (vlist:string list) =
                match vlist with
                | [] -> seq {yield valsSet}
                | v::rest -> 
                    let bkt = solution.Buckets.[v]
                    let lookupKey = valsSet |> Set.filter (fun (k,v) -> bkt.CScope.Contains k)
                    let cVals = bkt.Constraints.[lookupKey]
                    seq {for c in cVals do yield! loop (valsSet |> Set.add (v,c)) rest}
        yield! loop Set.empty solution.Variables
    }

let toSymbols = Set.map (fun (k:string,v) -> k,Symbols.Default.Symbol(v))
            

